# MIT License
#
# Copyright (c) 2023 Dechin CHEN
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import os
os.environ['GLOG_v'] = '4'
from quaternion import Quaternion
import numpy as np
import mindspore as ms
from mindspore import Tensor

if __name__ == '__main__':
    element = Tensor([0.], ms.float32)
    element_quaternion = Quaternion(element)
    print ('The type of element is: {}'.format(type(element_quaternion)))
    print ('The value of element is: {}'.format(element_quaternion.to_tensor()))
    assert np.allclose(element_quaternion.to_tensor().asnumpy(), np.array([[0., 0., 0., 0.]], np.float32))
    # The type of element is: <class 'quaternion.Quaternion'>
    # The value of element is: [[0. 0. 0. 0.]]

    vector = Tensor(np.arange(3), ms.float32)
    vector_quaternion = Quaternion(vector)
    print('The type of vector is: {}'.format(type(vector_quaternion)))
    print('The value of vector is: {}'.format(vector_quaternion.to_tensor()))
    assert np.allclose(vector_quaternion.to_tensor().asnumpy(), np.array([[0., 0., 1., 2.]], np.float32))
    # The type of vector is: <class 'quaternion.Quaternion'>
    # The value of vector is: [[0. 0. 1. 2.]]

    quater = Tensor(np.arange(4), ms.float32)
    quater_quaternion = Quaternion(quater)
    print('The type of quater is: {}'.format(type(quater_quaternion)))
    print('The value of quater is: {}'.format(quater_quaternion.to_tensor()))
    assert np.allclose(quater_quaternion.to_tensor().asnumpy(), np.array([[0., 1., 2., 3.]], np.float32))
    # The type of quater is: <class 'quaternion.Quaternion'>
    # The value of quater is: [[0. 1. 2. 3.]]

    batch_vec = Tensor(np.arange(12).reshape((4, 3)), ms.float32)
    batch_vec_quaternion = Quaternion(batch_vec)
    print('The type of batch_vec is: {}'.format(type(batch_vec_quaternion)))
    print('The value of batch_vec is: {}'.format(batch_vec_quaternion.to_tensor()))
    assert np.allclose(batch_vec_quaternion.to_tensor().asnumpy(), np.array([[ 0.,  0.,  1.,  2.],
                                                                             [ 0.,  3.,  4.,  5.],
                                                                             [ 0.,  6.,  7.,  8.],
                                                                             [ 0.,  9., 10., 11.]], np.float32))
    # The type of batch_vec is: <class 'quaternion.Quaternion'>
    # The value of batch_vec is: [[ 0.  0.  1.  2.]
    #  [ 0.  3.  4.  5.]
    #  [ 0.  6.  7.  8.]
    #  [ 0.  9. 10. 11.]]

    batch_quater = Tensor(np.arange(16).reshape((4, 4)), ms.float32)
    batch_quater_quaternion = Quaternion(batch_quater)
    print('The type of batch_quater is: {}'.format(type(batch_quater_quaternion)))
    print('The value of batch_quater is: {}'.format(batch_quater_quaternion.to_tensor()))
    assert np.allclose(batch_quater_quaternion.to_tensor().asnumpy(), np.array([[ 0.,  1.,  2.,  3.],
                                                                                [ 4.,  5.,  6.,  7.],
                                                                                [ 8.,  9., 10., 11.],
                                                                                [12., 13., 14., 15.]], np.float32))
    # The type of batch_quater is: <class 'quaternion.Quaternion'>
    # The value of batch_quater is: [[ 0.  1.  2.  3.]
    #  [ 4.  5.  6.  7.]
    #  [ 8.  9. 10. 11.]
    #  [12. 13. 14. 15.]]